import librosa
import numpy as np
import matplotlib.pyplot as plt
import os
import sounddevice as sd
# 设置matplotlib支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
import threading
import time
from matplotlib.animation import FuncAnimation

# ================== 配置参数 ==================
NORMAL_AUDIO_PATH = 'data/test2.mp3'  # 正常电机声音样本（训练/标定用）
TEST_AUDIO_PATH = 'data/test.mp3'  # 待检测的电机声音（可正常或异常）
FRAME_DURATION = 0.46  # 每帧时长（秒）
HOP_RATIO = 0.5  # 帧间重叠比例
SAMPLE_RATE = 44100  # 统一重采样率（推荐≥22050，最好44100）
# 对每个特征计算统计阈值（均值±??倍标准差）
RATIO = 3.0  # 可调节阈值灵敏度
# =============================================
# 1. 从正常样本中学习特征分布（建立阈值）
# =============================================


def extract_features(y, sr, frame_duration=0.1, hop_ratio=0.5):
    """
    提取音频的多维特征序列（按帧）
    
    参数:
    y: 音频时间序列数据
    sr: 采样率
    frame_duration: 每帧时长（秒）
    hop_ratio: 帧间重叠比例
    
    返回:
    包含rms、zcr、centroid特征的字典
    """
    frame_length = int(frame_duration * sr)
    hop_length = int(hop_ratio * frame_length)

    # 分帧提取特征
    # RMS能量：反映声音的整体响度或强度
    rms = librosa.feature.rms(y=y,
                              frame_length=frame_length,
                              hop_length=hop_length)[0]
    # 过零率：反映声音的粗糙度或频率感
    zcr = librosa.feature.zero_crossing_rate(y,
                                             frame_length=frame_length,
                                             hop_length=hop_length)[0]
    # 频谱质心：反映声音频谱的"明亮度"或高频能量占比
    centroid = librosa.feature.spectral_centroid(y=y,
                                                 sr=sr,
                                                 n_fft=frame_length,
                                                 hop_length=hop_length)[0]

    return {'rms': rms, 'zcr': zcr, 'centroid': centroid}


def train_thresholds(normal_audio_path, sr=SAMPLE_RATE):
    """
    使用正常音频训练特征阈值（基于统计学方法）
    
    参数:
    normal_audio_path: 正常音频文件路径
    sr: 采样率
    
    返回:
    各特征的阈值字典，包含均值、标准差和3σ区间
    """
    if not os.path.exists(normal_audio_path):
        raise FileNotFoundError(f"未找到正常样本: {normal_audio_path}")

    # 加载音频文件
    y, sr = librosa.load(normal_audio_path, sr=sr)
    # 提取特征
    features = extract_features(y, sr, FRAME_DURATION, HOP_RATIO)

    thresholds = {}

    for name, data in features.items():
        mean = np.mean(data)  # 计算均值
        std = np.std(data)  # 计算标准差
        upper = mean + RATIO * std  # 上限阈值
        lower = mean - RATIO * std  # 下限阈值
        thresholds[name] = {
            'mean': mean,
            'std': std,
            'upper': upper,
            'lower': lower
        }
        print(
            f"[训练] {name}: 均值={mean:.2f}, 标准差={std:.2f}, 阈值范围=[{lower:.2f}, {upper:.2f}]"
        )

    return thresholds


# =============================================
# 2. 对测试音频进行异常检测
# =============================================
def detect_anomaly(test_audio_path, thresholds, sr=SAMPLE_RATE):
    """
    检测测试音频是否异常
    
    参数:
    test_audio_path: 测试音频文件路径
    thresholds: 正常样本训练得到的阈值
    sr: 采样率
    
    返回:
    anomaly_flags: 每帧的异常状态数组
    time_axis: 时间轴
    features: 提取的特征
    """
    if not os.path.exists(test_audio_path):
        raise FileNotFoundError(f"未找到测试样本: {test_audio_path}")

    # 加载测试音频
    y, sr = librosa.load(test_audio_path, sr=sr)
    # 提取特征
    features = extract_features(y, sr, FRAME_DURATION, HOP_RATIO)

    # 存储每帧的异常状态
    n_frames = len(next(iter(features.values())))  # 获取帧数
    anomaly_flags = np.zeros(n_frames)  # 0: 正常, 1: 异常

    print("\n[检测结果]")
    # 将帧索引转换为时间轴
    time_axis = librosa.frames_to_time(range(n_frames),
                                       sr=sr,
                                       hop_length=int(HOP_RATIO *
                                                      FRAME_DURATION * sr))

    # 逐帧检测异常
    for i in range(n_frames):
        is_anomaly = False
        reasons = []

        # 检查每个特征是否超出阈值范围
        for name in features.keys():
            val = features[name][i]
            thresh = thresholds[name]
            if val < thresh['lower'] or val > thresh['upper']:
                is_anomaly = True
                reasons.append(
                    f"{name}={val:.2f} (阈值[{thresh['lower']:.2f},{thresh['upper']:.2f}])"
                )

        if is_anomaly:
            anomaly_flags[i] = 1
            print(
                f"⚠️  帧 {i} (时间 {time_axis[i]:.2f}s) 异常：{', '.join(reasons)}")

    # 判断整体是否异常：只要有1帧异常就算异常
    total_anomalies = np.sum(anomaly_flags)
    if total_anomalies > 0:
        print(f"\n❌ 检测到异常！共 {int(total_anomalies)} 个异常帧。")
    else:
        print(f"\n✅ 所有帧均在正常范围内，设备状态正常。")

    return anomaly_flags, time_axis, features


# =============================================
# 3. 可视化结果
# =============================================
def plot_results(features, anomaly_flags, time_axis, thresholds):
    """
    绘制特征曲线与异常点
    
    参数:
    features: 提取的特征数据
    anomaly_flags: 异常标记数组
    time_axis: 时间轴
    thresholds: 阈值数据
    """
    fig, axs = plt.subplots(3, 1, figsize=(14, 8), sharex=True)

    feature_names = ['rms', 'zcr', 'centroid']
    titles = ['RMS 能量', '过零率', '频谱质心 (Hz)']
    '''
    1. RMS 能量（Root Mean Square Energy）
    中文名：均方根能量 / 有效能量
    反映什么：声音的整体响度或强度
    🔍 物理意义：
    RMS 越高 → 声音越"响"或"强"
    对于电机：RMS 异常升高可能表示负载过大、振动加剧；异常降低可能表示停转或动力不足
    ✅ 应用场景（电机）：
    检测电机是否启动/停止
    判断运行负载是否变化
    发现异常振动或冲击噪声
    
    2. 过零率（Zero-Crossing Rate, ZCR）
    中文名：过零率
    反映什么：声音的粗糙度或频率感，粗略估计"像噪声还是像周期信号"
    🔍 物理意义：
    ZCR 高 → 信号变化快 → 像噪声、摩擦声、高频嘶嘶声
    ZCR 低 → 信号变化慢 → 像周期性、平稳的嗡嗡声（如正常电机）
    ✅ 应用场景（电机）：
    正常电机：低 ZCR（稳定电磁嗡鸣）
    故障电机：高 ZCR（如轴承磨损产生"咔哒"声、扫膛产生摩擦声）
    可用于检测"异响"、"干磨"、"松动"等故障
    
    3. 频谱质心（Spectral Centroid）
    中文名：频谱质心 / 频率重心
    反映什么：声音的频谱"明亮度"或"高频能量占比"
    🔍 物理意义：
    频谱质心高 → 高频成分多 → 声音"尖锐"、"刺耳"
    频谱质心低 → 低频成分多 → 声音"低沉"、"浑厚"
    ✅ 应用场景（电机）：
    正常电机：频谱集中在工频及其谐波（如 50Hz, 100Hz, 150Hz），质心较低
    故障电机（如轴承损坏）：产生高频冲击（2kHz~10kHz），频谱质心显著升高
    特别适合检测早期轴承故障，因为高频冲击在人耳不明显，但算法可捕捉
    '''
    colors = ['blue', 'green', 'purple']

    # 绘制每个特征的曲线
    for idx, (name, title,
              color) in enumerate(zip(feature_names, titles, colors)):
        ax = axs[idx]
        data = features[name]

        # 绘制正常点
        ax.plot(time_axis, data, label='正常', color=color, alpha=0.7)
        # 标记异常点为红色
        anomalous_time = time_axis[anomaly_flags == 1]
        anomalous_data = data[anomaly_flags == 1]
        ax.scatter(anomalous_time,
                   anomalous_data,
                   color='red',
                   s=30,
                   label='异常',
                   zorder=5)

        # 添加阈值线和正常范围区域
        thresh = thresholds[name]
        ax.axhline(thresh['upper'],
                   color='orange',
                   linestyle='--',
                   alpha=0.7,
                   linewidth=1)
        ax.axhline(thresh['lower'],
                   color='orange',
                   linestyle='--',
                   alpha=0.7,
                   linewidth=1)
        ax.fill_between(time_axis,
                        thresh['lower'],
                        thresh['upper'],
                        color='green',
                        alpha=0.1)

        ax.set_ylabel(title)
        ax.legend()
        ax.grid(True, alpha=0.3)

    axs[2].set_xlabel('时间 (秒)')
    plt.suptitle('电机声音特征与异常检测结果')
    plt.tight_layout()
    plt.show()


# =============================================
# 4. 实时播放和检测功能
# =============================================
def real_time_detection(test_audio_path, thresholds, sr=SAMPLE_RATE):
    """
    实时播放音频并动态检测异常
    
    参数:
    test_audio_path: 测试音频文件路径
    thresholds: 正常样本训练得到的阈值
    sr: 采样率
    """
    if not os.path.exists(test_audio_path):
        raise FileNotFoundError(f"未找到测试样本: {test_audio_path}")

    y, sr = librosa.load(test_audio_path, sr=sr)

    # 创建图形界面用于动态更新
    fig, axs = plt.subplots(3, 1, figsize=(14, 8), sharex=True)
    fig.suptitle('电机声音特征与异常检测结果（实时）')

    # 初始化数据存储
    feature_names = ['rms', 'zcr', 'centroid']
    titles = ['RMS 能量', '过零率', '频谱质心 (Hz)']
    colors = ['blue', 'green', 'purple']

    # 存储已经处理的数据 (使用线程安全的方式)
    processed_data = {name: [] for name in feature_names}
    processed_time = []
    anomaly_points = {name: {'time': [], 'data': []} for name in feature_names}

    # 线程同步机制
    data_lock = threading.Lock()
    update_event = threading.Event()

    frame_length = int(FRAME_DURATION * sr)
    hop_length = int(HOP_RATIO * frame_length)

    # 计算每帧对应的时间
    frames_total = len(y) // hop_length
    time_per_frame = FRAME_DURATION * HOP_RATIO

    # 用于线程间同步的变量
    current_frame = 0
    is_playing = True
    play_lock = threading.Lock()

    def audio_player():
        """
        音频播放线程函数
        负责播放音频并逐帧分析特征
        """
        nonlocal current_frame, is_playing
        print("\n🔊 开始实时播放和检测...")

        # 播放整个音频，同时进行分帧检测
        sd.play(y, sr)
        start_time = time.time()

        # 分帧处理但不逐帧等待播放
        for i in range(frames_total):
            with play_lock:
                if not is_playing:
                    break
                current_frame = i

            # 计算当前帧的起始和结束位置
            start_idx = i * hop_length
            end_idx = min(start_idx + frame_length, len(y))
            frame = y[start_idx:end_idx]

            # 提取当前帧的特征
            if len(frame) >= frame_length:
                # 提取特征
                rms = librosa.feature.rms(y=frame,
                                          frame_length=frame_length,
                                          hop_length=frame_length)[0]
                zcr = librosa.feature.zero_crossing_rate(
                    y=frame,
                    frame_length=frame_length,
                    hop_length=frame_length)[0]
                centroid = librosa.feature.spectral_centroid(
                    y=frame,
                    sr=sr,
                    n_fft=frame_length,
                    hop_length=frame_length)[0]

                current_features = {
                    'rms': rms[0],
                    'zcr': zcr[0],
                    'centroid': centroid[0]
                }
                current_time = i * time_per_frame

                # 检查是否异常
                is_anomaly = False
                reasons = []

                for name in feature_names:
                    val = current_features[name]
                    thresh = thresholds[name]

                    # 线程安全地更新数据
                    with data_lock:
                        processed_data[name].append(val)

                    if val < thresh['lower'] or val > thresh['upper']:
                        is_anomaly = True
                        reasons.append(f"{name}={val:.2f}")
                        with data_lock:
                            anomaly_points[name]['time'].append(current_time)
                            anomaly_points[name]['data'].append(val)

                with data_lock:
                    processed_time.append(current_time)

                if is_anomaly:
                    print(
                        f"⚠️  时间 {current_time:.2f}s 异常：{', '.join(reasons)}")
                else:
                    print(f"✅ 时间 {current_time:.2f}s 正常")

                # 触发更新事件
                update_event.set()

                # 控制处理速度与播放速度同步
                expected_time = start_time + current_time
                actual_time = time.time()
                if actual_time < expected_time:
                    time.sleep(expected_time - actual_time)

        with play_lock:
            is_playing = False
        # 最后触发一次更新以确保所有数据都被绘制
        update_event.set()

    # 动画更新函数
    def update_plot(frame):
        """
        图表更新函数，用于实时绘制检测结果
        """
        # 检查是否有新数据需要绘制
        if not update_event.is_set():
            return []
        update_event.clear()

        # 创建临时副本以避免在绘图时数据被修改
        with data_lock:
            time_copy = processed_time.copy()
            data_copy = {
                name: processed_data[name].copy()
                for name in feature_names
            }
            anomaly_copy = {
                name: {
                    'time': anomaly_points[name]['time'].copy(),
                    'data': anomaly_points[name]['data'].copy()
                }
                for name in feature_names
            }

        # 更新图表
        artists = []
        for idx, (name, title,
                  color) in enumerate(zip(feature_names, titles, colors)):
            ax = axs[idx]
            # 清除之前的内容
            ax.clear()

            # 只显示最近10秒的数据
            if time_copy:
                current_time = time_copy[-1] if time_copy else 0
                start_time = max(0, current_time - 10)  # 10秒窗口

                # 过滤数据，只保留最近10秒的
                filtered_indices = [
                    i for i, t in enumerate(time_copy) if t >= start_time
                ]
                if filtered_indices:
                    filtered_time = [time_copy[i] for i in filtered_indices]
                    filtered_data = [
                        data_copy[name][i] for i in filtered_indices
                    ]

                    # 绘制正常数据
                    line, = ax.plot(filtered_time,
                                    filtered_data,
                                    color=color,
                                    alpha=0.7,
                                    label='正常')
                    artists.append(line)

                    # 过滤异常点数据
                    anomaly_indices = [
                        i for i, t in enumerate(anomaly_copy[name]['time'])
                        if t >= start_time
                    ]
                    if anomaly_indices:
                        filtered_anomaly_time = [
                            anomaly_copy[name]['time'][i]
                            for i in anomaly_indices
                        ]
                        filtered_anomaly_data = [
                            anomaly_copy[name]['data'][i]
                            for i in anomaly_indices
                        ]

                        # 绘制异常点
                        scatter = ax.scatter(filtered_anomaly_time,
                                             filtered_anomaly_data,
                                             color='red',
                                             s=30,
                                             label='异常',
                                             zorder=5)
                        artists.append(scatter)

                # 设置x轴范围为最近10秒
                ax.set_xlim(max(0, current_time - 10), max(10, current_time))

            # 添加阈值线
            thresh = thresholds[name]
            line1 = ax.axhline(thresh['upper'],
                               color='orange',
                               linestyle='--',
                               alpha=0.7,
                               linewidth=1)
            line2 = ax.axhline(thresh['lower'],
                               color='orange',
                               linestyle='--',
                               alpha=0.7,
                               linewidth=1)
            artists.extend([line1, line2])

            fill = ax.fill_between([
                max(0, (time_copy[-1] if time_copy else 0) - 10),
                max(10, time_copy[-1] if time_copy else 10)
            ],
                                   thresh['lower'],
                                   thresh['upper'],
                                   color='green',
                                   alpha=0.1)
            artists.append(fill)

            ax.set_ylabel(title)
            ax.legend()
            ax.grid(True, alpha=0.3)

        axs[2].set_xlabel('时间 (秒)')
        plt.tight_layout()
        return artists

    # 创建并启动音频播放线程
    player_thread = threading.Thread(target=audio_player)
    player_thread.start()

    # 使用动画来更新图表
    ani = FuncAnimation(fig,
                        update_plot,
                        interval=100,
                        blit=False,
                        repeat=True)

    # 显示图形界面
    plt.show()

    # 等待播放线程完成
    player_thread.join()

    print("✅ 实时播放和检测完成")


# =============================================
# 5. 主程序
# =============================================
if __name__ == "__main__":
    print("🔊 电机声音异常检测系统（阈值法）\n")

    # Step 1: 训练阈值（使用正常样本）
    print("🔧 正在从正常样本中学习特征分布...")
    thresholds = train_thresholds(NORMAL_AUDIO_PATH)

    # 询问用户选择模式
    mode = input("\n请选择模式：\n1. 静态检测\n2. 实时播放检测\n请输入选项（1或2）：")

    if mode == "2":
        # 实时播放检测模式
        print(f"\n🔍 正在实时检测音频: {TEST_AUDIO_PATH}")
        real_time_detection(TEST_AUDIO_PATH, thresholds)
    else:
        # Step 2: 检测新样本
        print(f"\n🔍 正在检测音频: {TEST_AUDIO_PATH}")
        anomaly_flags, time_axis, features = detect_anomaly(
            TEST_AUDIO_PATH, thresholds)

        # Step 3: 可视化
        plot_results(features, anomaly_flags, time_axis, thresholds)
